package com.bigdata.hdfs;

import org.apache.commons.io.IOUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.*;

import java.io.BufferedWriter;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.nio.charset.Charset;
import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * hdfs文件操作
 *
 * @author chengye.yu
 */
public class HdfsService {
    private Log logger = LogFactory.getLog(HdfsService.class);
    private Configuration conf = null;

    public HdfsService(Configuration conf) {
        this.conf = conf;
    }

    /**
     * 获取HDFS文件系统
     *
     * @return org.apache.hadoop.fs.FileSystem
     */
    private FileSystem getFileSystem() throws IOException {
        return FileSystem.get(conf);
    }

    /**
     * 创建HDFS目录
     *
     * @param path HDFS的相对目录路径，比如：/testDir
     * @return boolean 是否创建成功
     * @throws Exception
     */
    public boolean mkdir(String path) throws Exception {
        // 如果目录已经存在，则直接返回
        if (checkExists(path)) {
            return true;
        } else {
            FileSystem fileSystem = null;

            try {
                fileSystem = getFileSystem();
                // 最终的HDFS文件目录
                String hdfsPath = generateHdfsPath(path);
                // 创建目录
                return mkdir(fileSystem, hdfsPath);
            } catch (Exception e) {
                logger.error(MessageFormat.format("创建HDFS目录失败，path:{0}", path), e);
                throw e;
            } finally {
                close(fileSystem);
            }
        }

    }

    /**
     * 上传文件至HDFS
     *
     * @param srcFile 本地文件路径，比如：D:/test.txt
     * @param dstPath HDFS的相对目录路径，比如：/testDir
     */
    public void uploadFileToHdfs(String srcFile, String dstPath) {
        this.uploadFileToHdfs(false, true, srcFile, dstPath);
    }

    /**
     * 上传文件至HDFS
     *
     * @param delSrc    是否删除本地文件
     * @param overwrite 是否覆盖HDFS上面的文件
     * @param srcFile   本地文件路径，比如：D:/test.txt
     * @param dstPath   HDFS的相对目录路径，比如：/testDir
     */
    public void uploadFileToHdfs(boolean delSrc, boolean overwrite, String srcFile, String dstPath) {
        FileSystem fileSystem = null;
        try {
            fileSystem = getFileSystem();
            //先判断文件夹是否存在
            boolean flag = this.mkdir(fileSystem, dstPath);
            if (flag) {
                // 源文件路径
                Path localSrcPath = new Path(srcFile);
                // 目标文件路径
                Path hdfsDstPath = new Path(generateHdfsPath(dstPath));
                fileSystem.copyFromLocalFile(delSrc, overwrite, localSrcPath, hdfsDstPath);
            }

        } catch (Exception e) {
            logger.error(MessageFormat.format("上传文件至HDFS失败，srcFile:{0},dstPath:{1}", srcFile, dstPath), e);
            throw new RuntimeException(MessageFormat.format("上传文件至HDFS失败，srcFile:{0},dstPath:{1}", srcFile, dstPath), e);
        } finally {
            close(fileSystem);
        }
    }

    /**
     * 不关闭fileSystem，在调用方关闭
     *
     * @param fileSystem
     * @param path
     * @return
     * @throws Exception
     */
    private boolean mkdir(FileSystem fileSystem, String path) throws Exception {
        // 如果目录已经存在，则直接返回
        if (checkExists(fileSystem, path)) {
            return true;
        } else {
            // 最终的HDFS文件目录
            String hdfsPath = generateHdfsPath(path);
            // 创建目录
            return fileSystem.mkdirs(new Path(hdfsPath));
        }

    }

    /**
     * 判断文件或者目录是否在HDFS上面存在
     * 不关闭fileSystem，在调用方关闭
     *
     * @param path HDFS的相对目录路径，比如：/testDir、/testDir/a.txt
     * @return boolean
     */
    private boolean checkExists(FileSystem fileSystem, String path) throws Exception {
        String hdfsPath = generateHdfsPath(path);
        return fileSystem.exists(new Path(hdfsPath));
    }

    /**
     * 判断文件或者目录是否在HDFS上面存在
     *
     * @param path HDFS的相对目录路径，比如：/testDir、/testDir/a.txt
     * @return boolean
     */
    public boolean checkExists(String path) {
        FileSystem fileSystem = null;
        try {
            fileSystem = getFileSystem();
            // 最终的HDFS文件目录
            String hdfsPath = generateHdfsPath(path);
            // 创建目录
            return checkExists(fileSystem, hdfsPath);
        } catch (Exception e) {
            logger.error(MessageFormat.format("'判断文件或者目录是否在HDFS上面存在'失败，path:{0}", path), e);
            return false;
        } finally {
            close(fileSystem);
        }
    }

    /**
     * 获取HDFS上面的某个路径下面的所有文件或目录（不包含子目录）信息
     *
     * @param path HDFS的相对目录路径，比如：/testDir
     * @return java.util.List<java.util.Map < java.lang.String, java.lang.Object>>
     */
    public List<Map<String, Object>> listFiles(String path, PathFilter pathFilter) {
        // 返回数据
        List<Map<String, Object>> result = new ArrayList<>();

        // 如果目录已经存在，则继续操作
        if (checkExists(path)) {
            FileSystem fileSystem = null;

            try {
                fileSystem = getFileSystem();

                // 最终的HDFS文件目录
                String hdfsPath = generateHdfsPath(path);

                FileStatus[] statuses;
                // 根据Path过滤器查询
                if (pathFilter != null) {
                    statuses = fileSystem.listStatus(new Path(hdfsPath), pathFilter);
                } else {
                    statuses = fileSystem.listStatus(new Path(hdfsPath));
                }

                if (statuses != null) {
                    for (FileStatus status : statuses) {
                        // 每个文件的属性
                        Map<String, Object> fileMap = new HashMap<>(2);

                        fileMap.put("path", status.getPath().toString());
                        // fileMap.put("isDir",status.isDirectory());

                        result.add(fileMap);
                    }
                }
            } catch (Exception e) {
                logger.error(MessageFormat.format("获取HDFS上面的某个路径下面的所有文件失败，path:{0}", path), e);
                throw new RuntimeException(MessageFormat.format("获取HDFS上面的某个路径下面的所有文件失败，path:{0}", path));
            } finally {
                close(fileSystem);
            }
        }

        return result;
    }

    /**
     * 从HDFS下载文件至本地
     *
     * @param srcFile HDFS的相对目录路径，比如：/testDir/a.txt
     * @param dstFile 下载之后本地文件路径（如果本地文件目录不存在，则会自动创建），比如：D:/test.txt
     */
    public void downloadFileFromHdfs(String srcFile, String dstFile) {
        // HDFS文件路径
        Path hdfsSrcPath = new Path(generateHdfsPath(srcFile));
        // 下载之后本地文件路径
        Path localDstPath = new Path(dstFile);

        FileSystem fileSystem = null;
        try {
            fileSystem = getFileSystem();

            fileSystem.copyToLocalFile(hdfsSrcPath, localDstPath);
        } catch (Exception e) {
            logger.error(MessageFormat.format("从HDFS下载文件至本地失败，srcFile:{0},dstFile:{1}", srcFile, dstFile), e);
            throw new RuntimeException(
                    MessageFormat.format("从HDFS下载文件至本地失败，srcFile:{0},dstFile:{1}", srcFile, dstFile));
        } finally {
            close(fileSystem);
        }
    }

    /**
     * 打开HDFS上面的文件并返回 InputStream
     *
     * @param path HDFS的相对目录路径，比如：/testDir/c.txt
     * @return FSDataInputStream
     */
    public FSDataInputStream open(String path) {
        // HDFS文件路径
        Path hdfsPath = new Path(generateHdfsPath(path));

        FileSystem fileSystem = null;
        try {
            fileSystem = getFileSystem();

            return fileSystem.open(hdfsPath);
        } catch (Exception e) {
            logger.error(MessageFormat.format("打开HDFS上面的文件失败，path:{0}", path), e);
            throw new RuntimeException(MessageFormat.format("打开HDFS上面的文件失败，path:{0}", path));
        }
    }

    /**
     * 打开HDFS上面的文件并返回byte数组，方便Web端下载文件
     * <p>
     * new ResponseEntity<byte[]>(byte数组, headers, HttpStatus.CREATED);
     * </p>
     * <p>
     * 或者：new ResponseEntity<byte[]>(FileUtils.readFileToByteArray(templateFile),
     * headers, HttpStatus.CREATED);
     * </p>
     *
     * @param path HDFS的相对目录路径，比如：/testDir/b.txt
     * @return FSDataInputStream
     */
    public byte[] openWithBytes(String path) {
        // HDFS文件路径
        Path hdfsPath = new Path(generateHdfsPath(path));

        FileSystem fileSystem = null;
        FSDataInputStream inputStream = null;
        try {
            fileSystem = getFileSystem();
            inputStream = fileSystem.open(hdfsPath);

            return IOUtils.toByteArray(inputStream);
        } catch (Exception e) {
            logger.error(MessageFormat.format("打开HDFS上面的文件失败，path:{0}", path), e);
            throw new RuntimeException(MessageFormat.format("打开HDFS上面的文件失败，path:{0}", path));
        } finally {
            if (inputStream != null) {
                try {
                    inputStream.close();
                } catch (Exception e) {
                    // ignore
                }
            }
        }

    }

    /**
     * 打开HDFS上面的文件并返回String字符串
     *
     * @param path HDFS的相对目录路径，比如：/testDir/b.txt
     * @return FSDataInputStream
     */
    public String openWithString(String path) {
        // HDFS文件路径
        Path hdfsPath = new Path(generateHdfsPath(path));

        FileSystem fileSystem = null;
        FSDataInputStream inputStream = null;
        try {
            fileSystem = getFileSystem();
            inputStream = fileSystem.open(hdfsPath);

            return IOUtils.toString(inputStream, Charset.forName("UTF-8"));
        } catch (Exception e) {
            logger.error(MessageFormat.format("打开HDFS上面的文件失败，path:{0}", path), e);
            throw new RuntimeException(MessageFormat.format("打开HDFS上面的文件失败，path:{0}", path));
        } finally {
            if (inputStream != null) {
                try {
                    inputStream.close();
                } catch (IOException e) {
                    // ignore
                }
            }
        }

    }

    /**
     * 重命名
     *
     * @param srcFile 重命名之前的HDFS的相对目录路径，比如：/testDir/b.txt
     * @param dstFile 重命名之后的HDFS的相对目录路径，比如：/testDir/b_new.txt
     */
    public boolean rename(String srcFile, String dstFile) {
        // HDFS文件路径
        Path srcFilePath = new Path(generateHdfsPath(srcFile));
        // 下载之后本地文件路径
        Path dstFilePath = new Path(dstFile);

        FileSystem fileSystem = null;
        try {
            fileSystem = getFileSystem();

            return fileSystem.rename(srcFilePath, dstFilePath);
        } catch (Exception e) {
            logger.error(MessageFormat.format("重命名失败，srcFile:{0},dstFile:{1}", srcFile, dstFile), e);
            throw new RuntimeException(MessageFormat.format("重命名失败，srcFile:{0},dstFile:{1}", srcFile, dstFile));
        } finally {
            close(fileSystem);
        }
    }

    /**
     * 删除HDFS文件或目录
     *
     * @param path HDFS的相对目录路径，比如：/testDir/c.txt
     * @return boolean
     */
    public boolean delete(String path) {
        // HDFS文件路径
        Path hdfsPath = new Path(generateHdfsPath(path));

        FileSystem fileSystem = null;
        try {
            fileSystem = getFileSystem();

            return fileSystem.delete(hdfsPath, true);
        } catch (Exception e) {
            logger.error(MessageFormat.format("删除HDFS文件或目录失败，path:{0}", path), e);
            throw new RuntimeException(MessageFormat.format("删除HDFS文件或目录失败，path:{0}", path));
        } finally {
            close(fileSystem);
        }

    }

    /**
     * 将字符串写入到文件的后面
     *
     * @param data
     * @param path
     */
    public void append(String data, String path) {
        FileSystem fileSystem = null;
        try {
            fileSystem = getFileSystem();
            // HDFS文件路径
            Path hdfsPath = new Path(generateHdfsPath(path));

            if (!fileSystem.exists(hdfsPath)) {
                fileSystem.createNewFile(hdfsPath);
            }
            FSDataOutputStream os = fileSystem.append(new Path(path));
            // 以UTF-8格式写入文件，不乱码
            BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(os, "utf-8"));
            writer.write(data);
            writer.write(System.lineSeparator());
            writer.close();
            os.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 获取某个文件在HDFS集群的位置
     *
     * @param path HDFS的相对目录路径，比如：/testDir/a.txt
     * @return org.apache.hadoop.fs.BlockLocation[]
     */
    public BlockLocation[] getFileBlockLocations(String path) {
        // HDFS文件路径
        Path hdfsPath = new Path(generateHdfsPath(path));

        FileSystem fileSystem = null;
        try {
            fileSystem = getFileSystem();
            FileStatus fileStatus = fileSystem.getFileStatus(hdfsPath);

            return fileSystem.getFileBlockLocations(fileStatus, 0, fileStatus.getLen());
        } catch (Exception e) {
            logger.error(MessageFormat.format("获取某个文件在HDFS集群的位置失败，path:{0}", path), e);
            throw new RuntimeException(MessageFormat.format("获取某个文件在HDFS集群的位置失败，path:{0}", path));
        } finally {
            close(fileSystem);
        }

    }

    /**
     * 将相对路径转化为HDFS文件路径
     *
     * @param dstPath 相对路径，比如：/data
     * @return java.lang.String
     */
    private String generateHdfsPath(String dstPath) {
        String hdfsPath = "";
        if (dstPath.startsWith("/")) {
            hdfsPath += dstPath;
        } else {
            hdfsPath = hdfsPath + "/" + dstPath;
        }

        return hdfsPath;
    }

    /**
     * close方法
     */
    private void close(FileSystem fileSystem) {
        if (fileSystem != null) {
            try {
                fileSystem.close();
            } catch (IOException e) {
                logger.error(e.getMessage());
            }
        }
    }

}
